MODULE fem_linear_solver

   USE fem_global_types
   USE fem_global_consts

   IMPLICIT NONE
   PRIVATE

   PUBLIC :: fem_init_linear_equations,       &
             fem_free_linear_equations,       &
             fem_solve_linear_equations

   REAL(REALK), SAVE, ALLOCATABLE :: B0(:)
   REAL(REALK), SAVE, ALLOCATABLE :: B(:)
   REAL(REALK), SAVE, ALLOCATABLE :: X(:)

CONTAINS

!-----------------------------------------------------------------------------

   SUBROUTINE fem_init_linear_equations(scheme,phi)

      USE fem_matrix_funcs, ONLY: fem_eliminate_BC

      IMPLICIT NONE
      TYPE(fem_scheme), INTENT(IN) :: scheme
      REAL(REALK),      INTENT(IN) :: phi(:,:,:)

      INTEGER(INTK) :: n(3)

      ! Only "interior" points
      ALLOCATE(B0(scheme%FEFNNB3D))
      ALLOCATE(B(scheme%FEFNNB3D))
      ALLOCATE(X(scheme%FEFNNB3D))

      ! Laplace RHS = 0
      B0(:) = zero
      ! FIXME put in switch for other equations RHS e.g. Poisson

      ! Eliminate boundary condition on phi from linear eqns
      n(:) = scheme%FEFNPt1D(:) - 2 
      CALL fem_eliminate_BC(scheme, n, phi, B0)
      B(:) = B0(:)

   END SUBROUTINE fem_init_linear_equations

!-----------------------------------------------------------------------------

   SUBROUTINE fem_free_linear_equations

      IMPLICIT NONE

      IF (ALLOCATED(B0)) DEALLOCATE(B0)
      IF (ALLOCATED(B)) DEALLOCATE(B)
      IF (ALLOCATED(X)) DEALLOCATE(X)

   END SUBROUTINE fem_free_linear_equations

!-----------------------------------------------------------------------------

   SUBROUTINE fem_solve_linear_equations(scheme,phi)

      USE fem_matrix_funcs, ONLY: fem_A_times_X,       &
                                  fem_push_phi_on_X,   &
                                  fem_push_X_on_phi

      IMPLICIT NONE
      TYPE(fem_scheme), INTENT(IN)    :: scheme
      REAL(REALK),      INTENT(INOUT) :: phi(:,:,:)

      INTEGER(INTK) :: n(3)
!integer :: i,j,k
!do k = 1, size(phi,3)
!do j = 1, size(phi,2)
!do i = 1, size(phi,1)
!   print *, i,j,k, phi(i,j,k)
!end do
!end do
!end do
      ! Initial guess using any known initial data points in interior
      X(:) = zero
      n(:) = scheme%FEFNPt1D(:) - 2 
!fixme      CALL fem_push_phi_on_X(n,phi,X)

      ! Call solver
      SELECT CASE ('PCG')
         CASE ('PCG'); CALL fem_get_phi_PCG(scheme,B,X)
!         CASE ('PCR  '); CALL Poi_Solv_PCR(B, X, Conv)
!         CASE ('JAC  '); CALL Poi_Solv_JAC(B, X, Conv)
!         CASE ('SOR  '); CALL Poi_Solv_SOR(B, X, Conv)
!         CASE ('SCHUR'); CALL Poi_Solv_Schur(B, X, Conv)
!         CASE DEFAULT ; STOP 'Laplace solver NYI'
      END SELECT

      ! Residual check
      CALL fem_A_times_X(scheme,n,X,B)
!do i = 1, size(B)
!   print '(I4, 3F10.5)', i, X(i), B(i), B0(i)
!enddo
      WRITE(IO1, *) 'Maximum residual =', MAXVAL(ABS(B0(:) - B(:)))
    
      ! Combine interior and boundary solutions
      CALL fem_push_X_on_phi(n,X,phi)

   END SUBROUTINE fem_solve_linear_equations

!-----------------------------------------------------------------------------

   SUBROUTINE fem_get_phi_PCG(scheme,B,X)
    
      USE fem_matrix_funcs, ONLY: fem_A_times_X

      IMPLICIT NONE
      TYPE(fem_scheme), INTENT(IN)    :: scheme
      REAL(REALK),      INTENT(IN)    :: B(:)
      REAL(REALK),      INTENT(INOUT) :: X(:)

      REAL(REALK), PARAMETER :: ConvThr = 1D-20
      INTEGER(INTK) :: iter, n3d, n(3)
      REAL(REALK)   :: B_Norm, Beta, Alpha, Res
      REAL(REALK)   :: P(scheme%FEFNNB3D)
      REAL(REALK)   :: Q(scheme%FEFNNB3D)
      REAL(REALK)   :: R(scheme%FEFNNB3D)
      REAL(REALK)   :: S(scheme%FEFNNB3D)

      REAL(REALK), EXTERNAL :: DDOT

      n3d = SIZE(B)
      n(:) = scheme%FEFNPt1D(:) - 2 
      B_Norm = DDOT(n3d, B, 1, B, 1)
      WRITE(IO1, *) ' B_Norm, ConvThr', B_Norm, ConvThr
      IF (B_norm < 1d-12) THEN
         X = zero
         RETURN
      END IF

      ! Initial RHS
      CALL fem_A_times_X(scheme,n,X,Q)

      ! Residual vector
      CALL DCOPY(n3d, B, 1, R, 1)
      CALL DAXPY(n3d, -One, Q, 1, R, 1)

      ! Pathway
      CALL Poi_Solv_PCon(scheme,R,S)
      Beta = Zero
      Beta = Beta + DDOT(n3d, R, 1, S, 1)
      Beta = One / Beta
      P(:) = Beta * S(:) 
!print *, "P:"
!print '(8F10.5)', P
!print *, "Q:"
!print '(8F10.5)', Q
!print *, "R:"
!print '(8F10.5)', R
!print *, "S:"
!print '(8F10.5)', S

      ! Iteration start
      WRITE(IO1,*) '(" ..... Enter    (PCG iteration   )")'
      DO iter = 1, n3d * 10
!print *, iter
!      DO iter = 1, 5

!print *, "fem_A_times_X..."
         CALL fem_A_times_X(scheme,n,P,Q)
!print *, "Q:"
!print '(8F10.5)', Q
         Alpha = Zero
         Alpha = Alpha + DDOT(n3d, P, 1, Q, 1)
         Alpha = One / Alpha

!print *, "DAXPY 1"
         CALL DAXPY(n3d,  Alpha, P, 1, X, 1)
!print *, "DAXPY 2"
         CALL DAXPY(n3d, -Alpha, Q, 1, R, 1)
!print *, "X:"
!print '(8F10.5)', X

         Res = DDOT(n3d, R, 1, R, 1) / B_Norm
         IF (MOD(iter,10) ==0) WRITE(IO1, *) 'PCG Iter.', iter, ' Res',Res
         IF (MOD(iter,10) ==0) print *, 'PCG Iter.', iter, ' Res',Res
         IF (Res < ConvThr) EXIT

!print *, "Poi_Solv_PCon"
         CALL Poi_Solv_PCon(scheme,R,S)

         Beta = Zero
         Beta = Beta + DDOT(n3d, R, 1, S, 1) 
         Beta = One / Beta
!print *, "DAXPY 3"
         CALL DAXPY(n3d, Beta, S, 1, P, 1)
    
      END DO

      WRITE(IO1, *) 'Iter.', iter, 'Res =', Res

   END SUBROUTINE fem_get_phi_PCG

!-----------------------------------------------------------------------------

   SUBROUTINE Poi_Solv_PCon(scheme,R,S)

      USE fem_matrix_funcs, ONLY: fem_get_Amat_inv

      IMPLICIT NONE
      TYPE(fem_scheme), INTENT(IN)  :: scheme
      REAL(REALK),      INTENT(IN)  :: R(:)
      REAL(REALK),      INTENT(OUT) :: S(:)

      REAL(REALK), POINTER :: Amat_dia_inv(:)

!     CASE ('Chole')
      CALL fem_get_Amat_inv(scheme,Amat_dia_inv)
      S(:) = Amat_Dia_Inv(:) * R(:)

!      SELECT CASE (GFCPreCond)
!         CASE ('None ')
!            OutFEF(:) = InpFEF(:)
!         CASE ('Point')
!            OutFEF(:) = DiaInvFEF(:) * InpFEF(:)
!         CASE ('Chole')
!            S(:) = Amat_Dia_Inv(:) * R(:)
!         CASE ('SOR ')
!            OutGTF(:) = DiaInvGTF(:) * InpGTF(:)
!            OutFEF(:) = DiaInvFEF(:) * InpFEF(:)
!            CALL Poi_Solv_SOR(InpGTF,InpFEF, OutGTF,OutFEF, .TRUE. , .TRUE. , GFCSORConv)
!      END SELECT

    END SUBROUTINE Poi_Solv_PCon

!-----------------------------------------------------------------------------

END MODULE fem_linear_solver
